#include <iostream>
#include <cuda_runtime.h>
#include "cublas_v2.h"
#include "../common/common.h"
#define DEVICENUM 0
#define M 2
#define K 3
#define N 4

int main(int argc, char **argv)
{
  cudaSetDevice(DEVICENUM);

  float alpha = 1.0f;
  float beta = 1.0f;

  float *A = (float *)malloc(sizeof(float) * M * K);
  float *B = (float *)malloc(sizeof(float) * K * N);
  float *C = (float *)malloc(sizeof(float) * M * N);
  float *res = (float *)malloc(sizeof(float) * M * N);

  initialDataIncreasing(A, M * K);
  initialDataIncreasing(B, K * N);
  memset(C, 0, sizeof(float) * M * N);
  memset(res, 0, sizeof(float) * M * N);

  printMatrix(A, M, K, "A");
  printMatrix(B, K, N, "B");
  // printMatrix(C, M, N, "C");

  float *dA, *dB, *dC;
  CHECK(cudaMalloc((void **)&dA, sizeof(float) * M * K));
  CHECK(cudaMalloc((void **)&dB, sizeof(float) * K * N));
  CHECK(cudaMalloc((void **)&dC, sizeof(float) * M * N));

  CHECK_CUBLAS(cublasSetMatrix(M, K, sizeof(float), A, M, dA, M)); // 逻辑上此时dA维度是[K,M]
  CHECK_CUBLAS(cublasSetMatrix(K, N, sizeof(float), B, K, dB, K)); // 逻辑上此时dB维度是[N,K]
  CHECK_CUBLAS(cublasSetMatrix(M, N, sizeof(float), C, M, dC, M)); // 逻辑上此时dC维度是[N,M]

  cublasHandle_t handle = 0; // 创建句柄
  CHECK_CUBLAS(cublasCreate(&handle));

  /* 4. A为行存储，B为行存储(alpha=beta=1.0f),此时dA和dB都要转置转置后计算结果如下
  矩阵dA_T:    矩阵dB_T:   sgemm计算结果:       C++按行读取后打印:
    1 2 3      1 2 3 4       38  44  50  56      38 83
    4 5 6      5 6 7 8       83  98 113 128      44 98
               9 10 11 125                       50 113
                                                 56 128
  */
  cudaMemset(dC, 0, sizeof(float) * M * N);
  CHECK_CUBLAS(cublasSgemm(handle,
                           CUBLAS_OP_T, // 需要对矩阵A转置
                           CUBLAS_OP_T, // 需要对矩阵B转置
                           M, N,        // 结果矩阵C的维度是[M,N]
                           K,           // 乘法左矩阵dA的列与乘法右矩阵dB的行=参数设置为K
                           &alpha,
                           dA, K, // dA矩阵，主维度此时为K。也可以记忆为设置CUBLAS_OP_T后矩阵A的列数
                           dB, N, // dB矩阵，主维度此时为N。也可以记忆为设置CUBLAS_OP_T后矩阵B的列数
                           &beta,
                           dC, M)); // dC矩阵，主维度此时为M
  CHECK_CUBLAS(cublasGetMatrix(M, N, sizeof(float), dC, M, res, M));
  printMatrix(res, N, M, "A&B row priority"); // 此时的结果是C_T

  free(A);
  free(B);
  free(C);
  free(res);
  cudaFree(dA);
  cudaFree(dB);
  cudaFree(dC);

  return 0;
}